package com.finance.cooperate.common.utils;

import org.dmg.pmml.FieldName;
import org.dmg.pmml.PMML;
import org.jpmml.evaluator.*;
import org.jpmml.model.PMMLUtil;
import org.springframework.core.io.ClassPathResource;
import org.xml.sax.SAXException;

import javax.xml.bind.JAXBException;
import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

/**
 * @ClassName ClassificationModel
 * @Description 模型加载工具类
 * @Author shen
 * @Date 2022/5/25 15:59
 * @Modify ...
 */
public class ClassificationModel {

    private Evaluator modelEvaluator;

    /**
     * @Author shen
     * @Description 通过传入 PMML 文件路径来生成机器学习模型
     * @Date 5:53 下午 2020/12/25
     * @Param [path pmml 文件路径]
     * @return
     **/
    public ClassificationModel(String path) {
        PMML pmml = null;

        try {

            ClassPathResource classPathResource = new ClassPathResource(path);
            InputStream is = classPathResource.getInputStream();

            if (is != null) {
                pmml = PMMLUtil.unmarshal(is);
                try {
                    is.close();
                } catch (IOException e) {
                    System.out.println("InputStream close error!");
                }
                ModelEvaluatorBuilder evaluatorBuilder = new ModelEvaluatorBuilder(pmml, (String) null)
                        .setModelEvaluatorFactory(ModelEvaluatorFactory.newInstance())
                        .setValueFactoryFactory(ValueFactoryFactory.newInstance());

                this.modelEvaluator = evaluatorBuilder.build();

                modelEvaluator.verify();
            }
        } catch (SAXException e) {
            e.printStackTrace();
        } catch (JAXBException e) {
            e.printStackTrace();
        } catch (IOException e) {
            e.printStackTrace();
        }

    }

    /**
     * @Author shen
     * @Description 获取模型需要的特征名称
     * @Date 5:54 下午 2020/12/25
     * @Param []
     * @return java.util.List<java.lang.String>
     **/
    public List<String> getFeatureNames() {
        List<String> featureNames = new ArrayList<String>();

        List<InputField> inputFields = modelEvaluator.getInputFields();

        for (InputField inputField : inputFields) {
            featureNames.add(inputField.getName().toString());
        }
        return featureNames;
    }


    /**
     * @Author shen
     * @Description 获取目标字段名称
     * @Date 5:54 下午 2020/12/25
     * @Param []
     * @return java.lang.String
     **/
    public String getTargetName() {

        return "label";

//         return modelEvaluator.getTargetFields().get(0).getName().toString();
    }

    /**
     * @Author shen
     * @Description 使用模型生成概率分布
     * @Date 7:02 下午 2020/12/25
     * @Param [arguments]
     * @return org.jpmml.evaluator.Classification
     **/
    private Object getProbabilityDistribution(Map<FieldName, ?> arguments) {
        Map<FieldName, ?> evaluateResult = modelEvaluator.evaluate(arguments);

        FieldName TargetName = FieldName.create(getTargetName());


        return evaluateResult.get(TargetName);

    }

    /**
     * @Author shen
     * @Description 预测标签值为1 的概率
     * @Date 5:55 下午 2020/12/25
     * @Param [arguments]
     * @return java.lang.Double
     **/
    public Double predictProba(Map<FieldName, Number> arguments) {


        Object distribution = getProbabilityDistribution(arguments);

        Double probability = null;


        if (distribution instanceof AffinityDistribution) {
            AffinityDistribution affinityDistribution = (AffinityDistribution) distribution;

        }


        if (probability == null) {
            ProbabilityDistribution probabilityDistribution = (ProbabilityDistribution) distribution;

            probability = probabilityDistribution.getProbability("1");
        }

        return probability;
    }

    /**
     * @Author shen
     * @Description 预测结果分类
     * @Date 5:56 下午 2020/12/25
     * @Param [arguments]
     * @return java.lang.Object
     **/
/*
    public Object predict(Map<FieldName, ?> arguments) {

        ProbabilityDistribution probabilityDistribution = getProbabilityDistribution(arguments);

        return probabilityDistribution.getPrediction();
    }
*/


}